package pipeline;

import two.phase.termination.TerminationToken;

import java.util.concurrent.*;

/**
 * Created by Administrator on 2017/3/17.
 */
public class ThreadPoolPipeDecorator<IN,OUT> implements Pipe<IN,OUT> {


    private final Pipe<IN,OUT> delegate;

    private final ExecutorService executorService;

    private final PipeTerminationToken terminationToken;

    private final CountDownLatch stageProcessDoneLatch = new CountDownLatch(1);

    public ThreadPoolPipeDecorator(Pipe<IN, OUT> delegate, ExecutorService executorService) {
        this.delegate = delegate;
        this.executorService = executorService;
        this.terminationToken = PipeTerminationToken.newInstance(executorService);
    }

    @Override
    public void init(PipeContext pipeContext) {
        delegate.init(pipeContext);
    }

    @Override
    public void process(final IN input) throws InterruptedException {
        Runnable task = new Runnable() {
            @Override
            public void run() {
                int remainingReservations = -1;
                try {
                    delegate.process(input);
                }catch (InterruptedException e){

                }finally {
                    remainingReservations = terminationToken.reservations.decrementAndGet();
                }

                if (terminationToken.isToShutdown() && 0 == remainingReservations ){
                    stageProcessDoneLatch.countDown();
                }
            }
        };

        executorService.submit(task);

        terminationToken.reservations.incrementAndGet();
    }

    @Override
    public void shutdown(long timeout, TimeUnit timeUnit) {
        terminationToken.setIsToShutdown();
        if (terminationToken.reservations.get() > 0){
            try{
                if (stageProcessDoneLatch.getCount() > 0){
                    stageProcessDoneLatch.await(timeout,timeUnit);
                }
            }catch (InterruptedException e) {

            }
        }

        delegate.shutdown(timeout,timeUnit);
    }

    @Override
    public void setNextPipe(Pipe<?, ?> nextPipe) {
        delegate.setNextPipe(nextPipe);
    }


    private static class PipeTerminationToken extends TerminationToken {
        private final static ConcurrentMap<ExecutorService,TerminationToken> INSTANCE_MAP = new ConcurrentHashMap<ExecutorService, TerminationToken>();

        private PipeTerminationToken() {
        }

        void setIsToShutdown(){
            this.toShutdown = true;
        }

        static TerminationToken newInstance(ExecutorService executorService){
            TerminationToken token = INSTANCE_MAP.get(executorService);
            if (null == token){
                token=new TerminationToken();
                TerminationToken existingToken = INSTANCE_MAP.putIfAbsent(executorService,token);
                if (null != existingToken){
                    token = existingToken;
                }
            }
            return token;
        }
    }
}
